package binary_heap;

import java.util.*;
import java.util.function.Function;

/**
 * 加强堆
 *
 * @param <T>
 */
@SuppressWarnings("unchecked")
public class PriorityQueueGreater<T> extends AbstractQueue<T> {

    private List<T> heap;

    private Map<Object, Integer> originIndexMap;

    private final Comparator<? super T> comparator;

    private final Function<T, Object> extractOriginFunc;

    public PriorityQueueGreater() {
        this(null, null);
    }

    public PriorityQueueGreater(Comparator<? super T> comparator) {
        this(comparator, null);
    }

    public PriorityQueueGreater(Comparator<? super T> comparator, Function<T, Object> extractOriginFunc) {
        this.heap = new ArrayList<>();
        this.comparator = comparator;
        this.extractOriginFunc = extractOriginFunc;
        if (null != extractOriginFunc) {
            this.originIndexMap = new HashMap<>();
        }
    }

    @Override
    public Iterator<T> iterator() {
        throw new UnsupportedOperationException();
    }

    @Override
    public int size() {
        return heap.size();
    }

    @Override
    public boolean offer(T t) {
        if (t == null) {
            throw new NullPointerException();
        }
        if (null != extractOriginFunc) {
            Object origin = extractOriginFunc.apply(t);
            Integer index = originIndexMap.get(origin);
            if (null != index) {
                heap.set(index, t);
                resign(heap, index);
                return true;
            }
        }

        heap.add(t);
        heapInsert(heap, heap.size() - 1, comparator);
        return true;
    }

    private void heapInsert(List<T> elements, int index, Comparator<? super T> comparator) {
        if (null != comparator) {
            heapInsertComparator(heap, index, comparator);
        } else {
            heapInsertComparable(heap, index);
        }
    }

    private void heapInsertComparable(List<T> elements, int index) {
        Comparable<T> current = ((Comparable<T>) elements.get(index));
        while (current.compareTo(elements.get(parentOf(index))) < 0) {
            swap(elements, index, parentOf(index));
            index = parentOf(index);
            current = (Comparable<T>) elements.get(index);
        }
    }

    private void heapInsertComparator(List<T> elements, int index, Comparator<? super T> comparator) {
        while (comparator.compare(elements.get(index), elements.get(parentOf(index))) < 0) {
            swap(elements, index, parentOf(index));
            index = parentOf(index);
        }
    }

    @Override
    public T poll() {
        T top = heap.get(0);

        int lastIndex = heap.size() - 1;
        swap(heap, 0, lastIndex);
        heap.remove(lastIndex);
        if (null != originIndexMap) {
            originIndexMap.remove(extractOriginFunc.apply(top));
        }

        heapify(heap, 0, comparator);
        return top;
    }

    private void heapify(List<T> elements, int index, Comparator<? super T> comparator) {
        if (null != comparator) {
            heapifyComparator(elements, index, comparator);
        } else {
            heapifyComparable(elements, index);
        }
    }

    private void heapifyComparable(List<T> elements, int index) {
        int left = leftOf(index);
        final int n = elements.size();
        while (left < n) {
            int right = left + 1;
            int smallerIndex = left;
            if (right < n && ((Comparable<T>) elements.get(smallerIndex)).compareTo(elements.get(right)) > 0) {
                smallerIndex = right;
            }
            smallerIndex = ((Comparable<T>) elements.get(index)).compareTo(elements.get(smallerIndex)) < 0
                    ? index : smallerIndex;
            if (index == smallerIndex) {
                break;
            }
            swap(elements, index, smallerIndex);
            index = smallerIndex;
            left = leftOf(index);
        }
    }

    private void heapifyComparator(List<T> elements, int index, Comparator<? super T> comparator) {
        int left = leftOf(index);
        final int n = elements.size();
        while (left < n) {
            int right = left + 1;
            int smallerChildIndex = left;
            if (right < n && comparator.compare(elements.get(smallerChildIndex), elements.get(right)) > 0) {
                smallerChildIndex = right;
            }
            if (comparator.compare(elements.get(index), elements.get(smallerChildIndex)) < 0) {
                break;
            }
            swap(elements, index, smallerChildIndex);
            index = smallerChildIndex;
            left = leftOf(index);
        }
    }

    @Override
    public T peek() {
        if (heap.isEmpty()) {
            throw new IndexOutOfBoundsException("is empty");
        }
        return heap.get(0);
    }

    private void resign(List<T> elements, int index) {
        heapInsert(elements, index, comparator);
        heapify(elements, index, comparator);
    }

    private static int parentOf(int index) {
        return (index - 1) / 2;
    }

    private static int leftOf(int index) {
        return index * 2 + 1;
    }

    private void swap(List<T> elements, int i, int j) {
        T a = elements.get(i);
        T b = elements.get(j);
        Collections.swap(elements, i, j);

        if (null != originIndexMap) {
            originIndexMap.put(extractOriginFunc.apply(a), j);
            originIndexMap.put(extractOriginFunc.apply(b), i);
        }
    }
}
